"""
边侧节点管理器
封装高级的边侧节点业务逻辑操作
"""

import asyncio
import logging
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
import traceback

from database import Node
from utils import LogManager
from exceptions import ValidationException
from .edge_client import EdgeClient
from .edge_types import (
    EdgeAggregationRequest, EdgeModelDeployment, EdgeBatchOperationResult
)
from process_manager import process_manager


class EdgeManager:
    """边侧节点管理器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.clients: Dict[str, EdgeClient] = {}
        self._active_tasks: Dict[str, str] = {}  # node_id -> task_id mapping
    
    def _get_client(self, node_id: str, endpoint: str) -> EdgeClient:
        """获取或创建边侧客户端"""
        if node_id not in self.clients:
            # EdgeClient不需要endpoint参数，会在具体方法调用时通过node参数确定URL
            self.clients[node_id] = EdgeClient()
        return self.clients[node_id]
    
    async def batch_health_check(self, nodes: List[Node]) -> Dict[str, bool]:
        """批量检查边侧节点健康状态 - 真正的并发处理"""
        if not nodes:
            return {}
        
        self.logger.info(f"开始批量健康检查 {len(nodes)} 个节点")
        
        # 创建所有健康检查任务
        tasks = []
        node_ids = []
        for node in nodes:
            task = asyncio.create_task(self._perform_health_check(node))
            tasks.append(task)
            node_ids.append(node.node_id)
            
            # 注册到进程管理器
            task_id = process_manager.register_async_task(
                task=task,
                name=f"health_check_{node.node_id}",
                task_type="health_check",
                related_nodes=[node.node_id]
            )
            self._active_tasks[node.node_id] = task_id
        
        try:
            # 并发执行所有任务，设置超时时间
            results = await asyncio.wait_for(
                asyncio.gather(*tasks, return_exceptions=True),
                timeout=10.0
            )
            
            # 处理结果
            health_status = {}
            for i, result in enumerate(results):
                node_id = node_ids[i]
                if isinstance(result, Exception):
                    self.logger.error(f"节点 {node_id} 健康检查异常: {str(result)}")
                    health_status[node_id] = False
                else:
                    _, is_healthy = result
                    health_status[node_id] = is_healthy
                
                # 清理任务注册
                if node_id in self._active_tasks:
                    process_manager.unregister_async_task(self._active_tasks[node_id])
                    del self._active_tasks[node_id]
            
            self.logger.info(f"批量健康检查完成: {sum(health_status.values())}/{len(health_status)} 节点在线")
            return health_status
            
        except asyncio.TimeoutError:
            self.logger.warning("批量健康检查超时，取消剩余任务")
            
            # 取消所有未完成的任务
            for task in tasks:
                if not task.done():
                    task.cancel()
            
            # 清理任务注册
            for node_id in node_ids:
                if node_id in self._active_tasks:
                    process_manager.unregister_async_task(self._active_tasks[node_id])
                    del self._active_tasks[node_id]
            
            # 返回所有节点为离线状态
            return {node.node_id: False for node in nodes}
        
        except Exception as e:
            self.logger.error(f"批量健康检查出现未预期错误: {str(e)}")
            
            # 清理任务注册
            for node_id in node_ids:
                if node_id in self._active_tasks:
                    process_manager.unregister_async_task(self._active_tasks[node_id])
                    del self._active_tasks[node_id]
            
            return {node.node_id: False for node in nodes}
    
    async def _perform_health_check(self, node: Node) -> Tuple[str, bool]:
        """对单个节点执行健康检查"""
        try:
            # 使用service_ip字段而非endpoint
            client = self._get_client(node.node_id, node.service_ip)
            health_status = await client.check_health(node)
            return node.node_id, health_status.status.value == "healthy"
        except Exception as e:
            self.logger.error(f"节点 {node.node_id} 健康检查失败: {str(e)}")
            return node.node_id, False
    
    async def batch_samples_info(self, db: Session, node_ids: List[str]) -> EdgeBatchOperationResult:
        """批量获取样本信息"""
        return await self._execute_batch_operation(
            db, node_ids, "samples_info", self._samples_info_single
        )
    
    async def batch_trigger_inference(self, db: Session, node_ids: List[str]) -> EdgeBatchOperationResult:
        """批量触发推理"""
        return await self._execute_batch_operation(
            db, node_ids, "trigger_inference", self._trigger_inference_single
        )
    
    async def batch_model_deployment(
        self, 
        db: Session, 
        node_ids: List[str], 
        model_id: str,
        deployment_config: Optional[Dict] = None
    ) -> EdgeBatchOperationResult:
        """批量模型部署"""
        async def deploy_single(node: Node) -> Dict[str, Any]:
            return await self._deploy_model_single(node, model_id, deployment_config)
        
        return await self._execute_batch_operation(
            db, node_ids, "deploy_model", deploy_single
        )
    
    async def coordinated_sample_aggregation(
        self,
        db: Session,
        node_ids: List[str],
        aggregation_params: Dict[str, Any]
    ) -> Dict[str, Any]:
        """协调的样本聚合操作"""
        try:
            nodes = await self._get_valid_edge_nodes(db, node_ids)
            
            # 构建聚合请求
            aggregation_request = EdgeAggregationRequest(
                difficulty_threshold=aggregation_params.get("difficulty_threshold"),
                time_start=aggregation_params.get("time_start"),
                time_end=aggregation_params.get("time_end"),
                max_samples=aggregation_params.get("max_samples", 1000)
            )
            
            results = {}
            total_samples = 0
            
            # 创建并发任务列表
            tasks = []
            for node in nodes:
                task = self._request_aggregation_with_timeout(node, aggregation_request)
                tasks.append((node.node_id, task))
            
            # 并发执行所有聚合请求
            if tasks:
                task_results = await asyncio.gather(
                    *[task for _, task in tasks], 
                    return_exceptions=True
                )
                
                # 处理结果
                for (node_id, _), result in zip(tasks, task_results):
                    if isinstance(result, Exception):
                        self.logger.error(f"节点 {node_id} 样本聚合失败: {str(result)}")
                        results[node_id] = {"success": False, "error": str(result)}
                    else:
                        results[node_id] = result.model_dump()
                        total_samples += result.sample_count
            
            return {
                "operation": "coordinated_sample_aggregation",
                "total_nodes": len(node_ids),
                "successful_nodes": [
                    node_id for node_id, result in results.items() 
                    if result.get("success", False)
                ],
                "total_aggregated_samples": total_samples,
                "results": results,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"协调样本聚合异常: {str(e)}")
            raise
    
    async def _request_aggregation_with_timeout(self, node: Node, request: EdgeAggregationRequest):
        """带超时的样本聚合请求"""
        try:
            # 设置超时时间为15秒
            return await asyncio.wait_for(
                self._get_client(node.node_id, node.service_ip).request_samples_aggregation(node, request),
                timeout=15.0
            )
        except asyncio.TimeoutError:
            raise Exception(f"节点 {node.node_id} 样本聚合请求超时")
        except Exception as e:
            raise e
    
    async def _execute_batch_operation(
        self,
        db: Session,
        node_ids: List[str],
        operation_name: str,
        operation_func
    ) -> EdgeBatchOperationResult:
        """执行批量操作的通用方法"""
        try:
            nodes = await self._get_valid_edge_nodes(db, node_ids)
            existing_node_ids = {node.node_id for node in nodes}
            missing_nodes = set(node_ids) - existing_node_ids
            
            results = {}
            successful_nodes = []
            failed_nodes = []
            
            # 创建并发任务列表
            tasks = []
            for node in nodes:
                # 为每个操作添加超时保护
                task = self._execute_single_with_timeout(operation_func, node, operation_name)
                tasks.append((node.node_id, task))
            
            # 并发执行所有操作
            if tasks:
                task_results = await asyncio.gather(
                    *[task for _, task in tasks], 
                    return_exceptions=True
                )
                
                # 处理结果
                for (node_id, _), result in zip(tasks, task_results):
                    if isinstance(result, Exception):
                        self.logger.error(f"节点 {node_id} 执行 {operation_name} 失败: {str(result)}")
                        results[node_id] = {"success": False, "error": str(result)}
                        failed_nodes.append(node_id)
                    else:
                        results[node_id] = result
                        successful_nodes.append(node_id)
            
            return EdgeBatchOperationResult(
                operation=operation_name,
                total_nodes=len(node_ids),
                successful_nodes=successful_nodes,
                failed_nodes=failed_nodes,
                missing_nodes=list(missing_nodes),
                results=results,
                timestamp=datetime.now().isoformat()
            )
            
        except Exception as e:
            self.logger.error(f"批量操作 {operation_name} 异常: {str(e)}")
            raise
    
    async def _execute_single_with_timeout(self, operation_func, node: Node, operation_name: str):
        """带超时的单节点操作执行"""
        try:
            # 根据操作类型设置不同的超时时间
            timeout_map = {
                "health_check": 10.0,      # 健康检查10秒超时
                "samples_info": 15.0,      # 样本信息15秒超时
                "trigger_inference": 30.0,  # 推理触发30秒超时
                "deploy_model": 60.0       # 模型部署60秒超时
            }
            timeout = timeout_map.get(operation_name, 15.0)
            
            return await asyncio.wait_for(
                operation_func(node),
                timeout=timeout
            )
        except asyncio.TimeoutError:
            raise Exception(f"节点 {node.node_id} 操作 {operation_name} 超时")
        except Exception as e:
            raise e
    
    async def _get_valid_edge_nodes(self, db: Session, node_ids: List[str]) -> List[Node]:
        """获取有效的边侧节点列表"""
        nodes = db.query(Node).filter(
            Node.node_id.in_(node_ids),
            Node.node_type == 'edge'
        ).all()
        
        if not nodes:
            raise ValidationException("没有找到有效的边侧节点")
        
        return nodes
    
    # 单节点操作方法
    async def _samples_info_single(self, node: Node) -> Dict[str, Any]:
        """单节点样本信息获取"""
        samples_info = await self._get_client(node.node_id, node.service_ip).get_samples_info(node)
        return {
            "success": True,
            "data": samples_info.model_dump()
        }
    
    async def _trigger_inference_single(self, node: Node) -> Dict[str, Any]:
        """单节点推理触发"""
        inference_result = await self._get_client(node.node_id, node.service_ip).trigger_inference(node)
        return inference_result.model_dump()
    
    async def _deploy_model_single(
        self, 
        node: Node, 
        model_id: str, 
        deployment_config: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """单节点模型部署"""
        from config import config
        
        # 构建模型部署配置
        cloud_base_url = config.cloud_server.get("public_url", "http://localhost:8000")
        model_download_url = f"{cloud_base_url}/api/v1/models/{model_id}/download"
        
        deployment = EdgeModelDeployment(
            model_id=model_id,
            model_url=model_download_url,
            model_path=deployment_config.get("model_path") if deployment_config else None
        )
        
        deployment_result = await self._get_client(node.node_id, node.service_ip).deploy_model(node, deployment)
        return deployment_result.model_dump()
    
    def cleanup_all_tasks(self):
        """清理所有活动任务"""
        self.logger.info("清理EdgeManager中的所有活动任务")
        
        # 取消所有边侧相关的异步任务
        edge_tasks = process_manager.get_edge_connection_tasks()
        for task_info in edge_tasks:
            if task_info.task and not task_info.task.done():
                task_info.cancel()
                self.logger.info(f"取消边侧任务: {task_info.task_id}")
        
        # 清理客户端连接
        for node_id, client in self.clients.items():
            try:
                # 如果客户端有清理方法，调用它
                if hasattr(client, 'cleanup'):
                    client.cleanup()
            except Exception as e:
                self.logger.warning(f"清理客户端 {node_id} 时出错: {str(e)}")
        
        self.clients.clear()
        self._active_tasks.clear()
        
        self.logger.info("EdgeManager任务清理完成")

# 全局边侧管理器实例
edge_manager = EdgeManager() 